import torch
import torch_sparse

# 假设我们有一个图，其中有5个节点
num_nodes = 5
# edge_index 表示边的起始和结束节点
# 这里我们有一些重复边，例如从节点0到节点1有两条边
total_edge_index = torch.tensor([[0, 0, 1, 2, 3, 4, 4],
                                 [1, 1, 1, 3, 4, 0, 0]])
# edge_value 表示每条边的权重
# 对于重复边，我们想要将它们的权重相加
total_edge_value = torch.tensor([1., 1., 1., 1., 1., 1., 1.])

# 使用 torch_sparse.coalesce 合并重复边
# op='add' 表示将重复边的权重相加
coalesced_edge_index, coalesced_edge_value = torch_sparse.coalesce(
    total_edge_index,
    total_edge_value,
    m=num_nodes,
    n=num_nodes, op='add')

print(coalesced_edge_index)
# 输出: tensor([[0, 0, 1, 2, 3, 4, 4],
#         [1, 2, 1, 3, 4, 0, 1]])
print(coalesced_edge_value)
# 输出: tensor([1, 2, 3, 4, 5, 6, 7])
